from accuracy import RMSE, MAE, MAPE, R_square
from model import *
from read_data import *
import time
from timeit import default_timer as timer
import numpy as np
import pandas as pd
import warnings
import math

from HHO_LSTM import HHO
warnings.filterwarnings("ignore")
# 自己预留的代码
warnings.filterwarnings("ignore")
# In[158]:


def run_hho_lstm(city, mod, sequence_length, horizon, decomposition):
    times = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time()))
    filename = str(times)+'_' + "h"+str(horizon)+"_"+mod + ".csv"
    figname = str(times)+'_' + "h"+str(horizon)+"_"+mod+".png"
    savepath = save_path+city+"/"+mod + "/"
    generpath(savepath)
    datapath = data_path + 'data'+"_" + decomposition + ".csv"

    if decomposition in mod:
     #   "IMF1","IMF2","IMF3","IMF4","IMF5","IMF6","IMF7","IMF8","IMF9","IMF10"
        # a = ["IMF4"]
        a = ["imf0", "imf1", 'imf2', "imf3", "imf4", "imf5", "imf6", "imf7"]
        for i in a:
            saveindex = i
            index = i
            horizon = horizon

            def function(variables_values=[32, 16]):
                x_train_initial, x_test_initial, y_train_initial, y_test_initial, x_train, x_test, y_train, y_test = read_data(
                    datapath, index, sequence_length, horizon)
                epochs = 50
                verbose = 0
                model = Sequential()
                model.add(LSTM(int((round(variables_values[0]))), input_shape=(
                    int(x_train.shape[1]), 1), return_sequences=True))
                model.add(LSTM(int((round(variables_values[1])))))
                model.add(Dense(1))
                model.compile(loss='mean_squared_error', optimizer='adam')
                model.fit(x_train, y_train, epochs=epochs,
                          validation_split=0, verbose=verbose)

                result = model.predict(x_test)  # x_test
                result = result.reshape(-1, 1)
                fuctions = RMSE(y_test, result)  # y_test
                print("REMS是", fuctions)
                print(fuctions)
                return fuctions

            s = HHO(objf=function,  dim=2,
                    SearchAgents_no=4,
                    lb=20,
                    ub=50,
                    Max_iter=3)
            # sequence_length = int(hho[2])
            # hho = hho[:2]
            print(s)
            ans = s.bestIndividual
            print(ans)
            print(s.best)
            x_train_initial, x_test_initial, y_train_initial, y_test_initial, x_train, x_test, y_train, y_test = read_data(
                datapath, index, sequence_length, horizon)
            # print("x_train", x_train.shape)
            # print("y_train", y_train.shape)
            # print("x_test", x_test.shape)
            # print("y_test", y_test.shape)
            verbose = 0
            n_multi = 0
            model = Sequential()
            model.add(LSTM(int(round(ans[0])),input_shape=(int(x_train.shape[1]),1),return_sequences=True))
            model.add(LSTM(int(round(ans[1]))))
            model.add(Dense(1))
            model.compile(loss='mean_squared_error',optimizer='adam')
            model.fit(x_train,y_train,epochs=epochs,validation_split=0,verbose=verbose)
            pre1 = model.predict(x_test)
            y_test_rel,y_test_pre = iverse_data(y_train,pre1,y_test)

            print("----------", i, "----------")
            print('MAE：', MAE(y_test_rel, y_test_pre))
            print('RMSE：', RMSE(y_test_rel, y_test_pre))
            print('MAPE：', MAPE(y_test_rel, y_test_pre)[0])
            print('R方：', R_square(y_test_rel, y_test_pre))
            generfile(savepath, filename, len(y_test))
            datasave(savepath+filename, saveindex, y_test_pre)
            print("----------", index, "had finished", "----------")
# In[]
        savepath1 = savepath+filename
        figpath = savepath + str(times)+"_"+mod+".png"
        datapr = pd.read_csv(savepath1)
        datapr1 = np.array(datapr)[:, 2:]
        y_pre = np.sum(datapr1, axis=1)
        y_pre = y_pre.reshape(-1, 1)
        datare = pd.read_csv(data_path+origin_data)
        datare.fillna(method='pad', inplace=True)
        y_rel = np.array(datare[city][-len(y_pre):]).reshape(-1, 1)
        y_pre = np.array(y_pre).reshape(len(y_pre),)
        y_rel = np.array(y_rel).reshape(len(y_pre),)

        generfile(savepath, str(times)+'_' + "合并结果.csv", len(y_test))
        datasave(savepath+str(times)+'_' + "合并结果.csv",
                 mod+"_h"+str(horizon), y_pre)

        preplot(y_pre, y_rel, figpath)

        hho_lstm_mae = MAE(y_rel, y_pre)
        hho_lstm_rmse = RMSE(y_rel, y_pre)
        hho_lstm_mape = MAPE(y_rel, y_pre)
        hho_lstm_R = R_square(y_rel, y_pre)

        hho_lstm_zb = [hho_lstm_mae, hho_lstm_rmse, hho_lstm_mape, hho_lstm_R]
        hho_lstm_zb = np.array(hho_lstm_zb).reshape(-1, 1)
        zbfile = str(times)+"_指标" + ".csv"
        generfile(savepath, zbfile, len(hho_lstm_zb))
        datasave(savepath+zbfile, "pso_lstm", hho_lstm_zb)

        print(city, mod, horizon)
        print('MAE：', hho_lstm_mae)
        print('RMSE：', hho_lstm_rmse)
        print('MAPE：', hho_lstm_mape)
        print('R方：',hho_lstm_R)

 # In[153]:
data_path = ""
save_path = "预测/Kmeans-alo-LSTM"
origin_data = "data.csv"
decomp_index = ['CEEMDAN8']
n_multi = 0
verbose = 0
sequence_length = 8
netmodel = "HHO_LSTM"
decomposition = decomp_index[0]
mod = netmodel + '_' + decomposition
cityindx = ["Power"]
epochs_index = [35]
hindex = [1]
# """
# In[153]:
for city in cityindx:
    for h in hindex:
        for epo in epochs_index:
            epochs = epo
            horizon = h
            city = city
            print("----------", city, mod, "h:", horizon,
                  "epochs:", epochs, "----------")
            run_hho_lstm(city, mod, sequence_length, horizon, decomposition)

# """
